"""
2019.7.4 计算旋转的iou
参考：https://blog.csdn.net/qq_29296685/article/details/99979069
此网站上还有旋转框的 NMS 函数
"""
# coding=utf-8
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import numpy as np
import cv2
import tensorflow as tf
import time
import torch


def rbbx_overlaps(boxes1, boxes2, gpu_id):
    pass


"""
计算旋转面积
boxes1,boxes2格式为x,y,w,h,theta
"""


def iou_rotate_calculate(boxes1, boxes2, use_gpu=False, gpu_id=1):
    #print("####boxes2:", boxes1.shape)
    #print("####boxes2:", boxes2.shape)
    #os._exit()
    # start = time.time()
    if use_gpu:
        print("暂时用不了")
        ious = rbbx_overlaps(boxes1, boxes2, gpu_id)
    else:
        area1 = boxes1[2] * boxes1[3]
        area2 = boxes2[2] * boxes2[3]
        r1 = ((boxes1[0], boxes1[1]), (boxes1[2], boxes1[3]), boxes1[4])
        r2 = ((boxes2[0], boxes2[1]), (boxes2[2], boxes2[3]), boxes2[4])
        int_pts = cv2.rotatedRectangleIntersection(r1, r2)[1]
        if int_pts is not None:
            order_pts = cv2.convexHull(int_pts, returnPoints=True)
            int_area = cv2.contourArea(order_pts)
            # 计算出iou
            ious = int_area * 1.0 / (area1 + area2 - int_area)
        else:
            ious = 0
    return ious


"""
计算旋转面积
boxes1,boxes2格式为x,y,w,h,theta
输出为张量格式
"""


def iou_rotate_calculate1(boxes1, boxes2, use_gpu=False, gpu_id=1):
    # 将gpu的tensor转成cpu的
    boxes1 = boxes1.cpu()
    boxes2 = boxes2.cpu()
    # 将tensor转成numpy的
    boxes1 = boxes1.numpy()
    boxes2 = boxes2.numpy()
    ious_total = []
    # start = time.time()
    if use_gpu:
        print("@@@@暂时用不了")
        ious = rbbx_overlaps(boxes1, boxes2, gpu_id)
    else:
        for num in range(0, len(boxes2)):
            area1 = boxes1[0] * boxes1[1]
            area2 = boxes2[num, 2] * boxes2[num, 3]
            r1 = ((boxes2[num, 0], boxes2[num, 1]), (boxes1[0], boxes1[1]), boxes1[2])
            r2 = ((boxes2[num, 0], boxes2[num, 1]), (boxes2[num, 2], boxes2[num, 3]), boxes2[num, 4])
            int_pts = cv2.rotatedRectangleIntersection(r1, r2)[1]
            if int_pts is not None:
                order_pts = cv2.convexHull(int_pts, returnPoints=True)
                int_area = cv2.contourArea(order_pts)
                # 计算出iou
                ious = int_area * 1.0 / (area1 + area2 - int_area)
                ious_total.append(ious)
            else:
                ious = 0
                ious_total.append(ious)
        # numpy转为CPU tensor
        total_ious = torch.from_numpy(np.array(ious_total))
        total_ious = total_ious.type(torch.FloatTensor)  # 转Float
        # CPU tensor转GPU tensor
        ious = total_ious.cuda()
    return ious


"""
计算旋转面积
boxes1,boxes2格式为x,y,w,h,theta
输出为张量格式
"""


def bbox_iou_rotate_calculate1(boxes1, boxes2, use_gpu=False, gpu_id=1):
    # 将gpu的tensor转成cpu的
    boxes1 = boxes1.cpu()
    boxes2 = boxes2.cpu()
    # 将tensor转成numpy的
    boxes1 = boxes1.numpy()
    boxes2 = boxes2.numpy()
    ious_total = []
    # start = time.time()
    if use_gpu:
        print("暂时用不了")
        ious = rbbx_overlaps(boxes1, boxes2, gpu_id)
    else:
        for num in range(0, len(boxes2)):
            area1 = boxes1[num, 2] * boxes1[num, 3]
            area2 = boxes2[num, 2] * boxes2[num, 3]
            r1 = ((boxes1[num, 0], boxes1[num, 1]), (boxes1[num, 2], boxes1[num, 3]), boxes1[num, 4])
            r2 = ((boxes2[num, 0], boxes2[num, 1]), (boxes2[num, 2], boxes2[num, 3]), boxes2[num, 4])
            int_pts = cv2.rotatedRectangleIntersection(r1, r2)[1]
            if int_pts is not None:
                order_pts = cv2.convexHull(int_pts, returnPoints=True)
                int_area = cv2.contourArea(order_pts)
                # 计算出iou
                ious = int_area * 1.0 / (area1 + area2 - int_area)
                ious_total.append(ious)
            else:
                ious = 0
                ious_total.append(ious)
        # numpy转为CPU tensor
        total_ious = torch.from_numpy(np.array(ious_total))
        total_ious = total_ious.type(torch.FloatTensor)  # 转Float
        # CPU tensor转GPU tensor
        ious = total_ious.cuda()
    return ious


if __name__ == '__main__':
    import os
    import numpy as np
    #os.environ["CUDA_VISIBLE_DEVICES"] = '1'

    theta = np.linspace(0,180,181)
    #theta = [20,30]
    h_base1 = 2
    h_base2 = 5
    ratio = np.arange(1,2,0.1)
    rious = np.empty((len(ratio),len(theta)))      # 用于保存IoU的值
    m = 0
    for i in ratio:
        n = 0
        for j in theta:
            pre_box = [0,0,h_base1*i,h_base1,0]
            tru_box = [0,0,h_base1*i,h_base1,j]
            boxes1 = np.array(pre_box, np.float32)

            boxes2 = np.array(tru_box, np.float32)

            ious_area = iou_rotate_calculate(boxes1, boxes2, use_gpu=False)
            rious[m,n] = ious_area
            n = n+1
            #print('theta:%.2f,ratio:%d,ious_area:%.2f'%(i,j,ious_area))
        m = m+1
    # 绘制不同长宽比下，随着角度变化Iou的变换情况
    import matplotlib.pyplot as plt
    cmap = plt.get_cmap('viridis')     # 选择一个colormap
    # colormap是一个函数，它可以获取从0到1的值数组，
    # 并将它们映射到rgba颜色。cmap生成长度从0到1的等距数字数组。因此，
    colors = cmap(np.linspace(0, 1, len(ratio)))

    plt.title('IoU relationship with angle and ratio of w/h')
    for index,(value,color) in enumerate(zip(ratio,colors)):
        plt.plot(theta, rious[index], color=color, label='ratio={:.2}'.format(value))

    plt.legend()  # 显示图例

    plt.xlabel('theta')
    plt.ylabel('IoU')
    plt.show()
pass




#python 一个折线图绘制多个曲线
